#include <iostream>
#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/Core>
#include <fstream>
#include <random>
#include <chrono>
#include <cmath>

class State
{
public:
    State() : x_(0), y_(0), theta_(0)
    {}

    double x() const { return x_; }
    double y() const { return y_; }
    double theta() const { return theta_; }

    double &x() { return x_; }
    double &y() { return y_; }
    double &theta() { return theta_; }

private:
    double x_;
    double y_;
    double theta_;
};

class Control
{
public:
    Control() : v_(0), dtheta_(0)
    {}

    double v() const { return v_; }
    double dtheta() const { return dtheta_; }

    double &v() { return v_; }
    double &dtheta() { return dtheta_; }

private:
    double v_;
    double dtheta_;
};

class Measurement
{
public:
    Measurement() : d1_(0), d2_(0), theta_(0)
    {}

    double d1() const { return d1_; }
    double d2() const { return d2_; }
    double theta() const { return theta_; }
    
    double& d1() { return d1_; }
    double& d2() { return d2_; }
    double& theta() { return theta_; }

private:
    double d1_;
    double d2_;
    double theta_;
};

class ActionModel
{
public:
    ActionModel()
    {
        cov_.setIdentity();
    }

    State transition(const State &x, const Control &u)
    {
        State x_t;
        x_t.theta() = x.theta() + u.dtheta();
        x_t.x() = x.x() + std::cos(x_t.theta()) * u.v();
        x_t.y() = x.y() + std::sin(x_t.theta()) * u.v();
        return x_t;
    }

    void updateJacobians(const State &x, const Control &u)
    {
        F_.setZero();
        F_(0, 0) = 1.0;
        F_(0, 1) = 0.0;
        F_(0, 2) = -std::sin(x.theta() + u.dtheta()) * u.v();

        F_(1, 0) = 0.0;
        F_(1, 1) = 1.0;
        F_(1, 2) = std::cos(x.theta() + u.dtheta()) * u.v();

        F_(2, 0) = 0.0;
        F_(2, 1) = 0.0;
        F_(2, 2) = 1.0;
    }

    const Eigen::Matrix3d& F() const
    {
        return F_;
    }

    const Eigen::Matrix3d& getCovariance() const
    {
        return cov_;
    }

private:
    Eigen::Matrix3d F_;
    Eigen::Matrix3d cov_;
};

class MeasureModel
{
public:
    MeasureModel(double landmark1x, double landmark1y, double landmark2x, double landmark2y)
    {
        lanemark1_ << landmark1x, landmark1y;
        lanemark2_ << landmark2x, landmark2y;
        cov_.setIdentity();
    }

    Measurement measure(const State& x) const
    {
        Measurement measurement;
        Eigen::Vector2d position(x.x(), x.y());
        
        Eigen::Vector2d delta1 = position - lanemark1_;
        measurement.d1() = std::sqrt( delta1.dot(delta1) );
        
        Eigen::Vector2d delta2 = position - lanemark2_;
        measurement.d2() = std::sqrt( delta2.dot(delta2) );

        measurement.theta() = x.theta();
        
        return measurement;
    }

    void updateJacobians( const State& x )
    {
        Eigen::Vector2d position(x.x(), x.y());
        
        Eigen::Vector2d delta1 = position - lanemark1_;
        Eigen::Vector2d delta2 = position - lanemark2_;
        
        double d1 = std::sqrt( delta1.dot(delta1) );
        double d2 = std::sqrt( delta2.dot(delta2) );

        H_.setZero();
        H_(0, 0) = delta1.x() / d1;
        H_(0, 1) = delta1.y() / d1;
        H_(0, 2) = 0.0;

        H_(1, 0) = delta2.x() / d2;
        H_(1, 1) = delta2.y() / d2;
        H_(1, 2) = 0.0;

        H_(2, 0) = 0.0;
        H_(2, 1) = 0.0;
        H_(2, 2) = 1.0;
    }

    const Eigen::Matrix3d& H() const
    {
        return H_;
    }

    const Eigen::Matrix3d& getCovariance() const
    {
        return cov_;
    }

private:
    Eigen::Vector2d lanemark1_;
    Eigen::Vector2d lanemark2_;
    Eigen::Matrix3d H_;
    Eigen::Matrix3d cov_;
};

class ExtendedKalmanFilter 
{
public:
    ExtendedKalmanFilter(State x0, Eigen::Matrix3d P0) : x_(x0), P_(P0)
    {
        
    }

    const State &predict(ActionModel &s, const Control& u)
    {
        s.updateJacobians(x_, u);
        x_ = s.transition(x_, u);
        P_ = (s.F() * P_ * s.F().transpose()) + s.getCovariance();
        return x_;
    }

    const State& update(MeasureModel& m, const Measurement& z)
    {
        m.updateJacobians(x_);
        Eigen::Matrix3d S = (m.H() * P_ * m.H().transpose()) + m.getCovariance();
        Eigen::Matrix3d K = P_ * m.H().transpose() * S.inverse();
        
        Eigen::Vector3d delta;
        delta.x() = z.d1() - m.measure(x_).d1();
        delta.y() = z.d2() - m.measure(x_).d2();
        delta.z() = z.theta() - m.measure(x_).theta();
        Eigen::Vector3d rise = K * delta;

        x_.x() = x_.x() + rise.x();
        x_.y() = x_.y() + rise.y();
        x_.theta() = x_.theta() + rise.z();

        P_ -= K * m.H() * P_;
        return x_;
    }

private:
    Eigen::Matrix3d P_;
    State x_;
};


int main0()
{
    State x; 
    Control u;
    ActionModel sys;
    MeasureModel pm(-10, -10, 30, 75);

    std::default_random_engine generator;
    generator.seed( std::chrono::system_clock::now().time_since_epoch().count() );
    std::normal_distribution<double> noise(0, 1);

    Eigen::Matrix3d Q = Eigen::Matrix3d::Identity();
    ExtendedKalmanFilter preditor(x, Q);
    ExtendedKalmanFilter ekf(x, Q);

    double systemNoise = 0.1;
    double orientationNoise = 0.025;
    double distanceNoise = 0.25;

    const size_t N = 100;
    std::ofstream ofs("kalman.csv");
    for(size_t i = 1; i <= N; i++)
    {
        u.v() = 1. + std::sin(2.0 * M_PI / N);
        if(i > 50)
            u.dtheta() = -std::sin(2.0 * M_PI / N);
        else
            u.dtheta() =  std::sin(2.0 * M_PI / N);
            
        auto x_predict = preditor.predict(sys, u);
        auto x_ekf = ekf.predict(sys, u);

        x = sys.transition(x, u);
        std::cout << "noise : " << noise(generator) << std::endl;
        x.x() += systemNoise * noise(generator);
        x.y() += systemNoise * noise(generator);
        x.theta() += systemNoise * noise(generator);
    
        Measurement position = pm.measure(x);
        position.d1() += distanceNoise * noise(generator);
        position.d2() += distanceNoise * noise(generator);
        position.theta() += orientationNoise * noise(generator);
        x_ekf = ekf.update(pm, position);     

        ofs << x.x() << "," << x.y() << "," << x.theta() << ","
                    << x_predict.x() << "," << x_predict.y() << "," << x_predict.theta() << ","
                    << x_ekf.x() << "," << x_ekf.y() << "," << x_ekf.theta()
                    << std::endl;   
    }

    return -1;
}